package com.dylan.redislimiting.aspect;

import com.dylan.redislimiting.annotation.LimitType;
import com.dylan.redislimiting.annotation.RateLimiting;
import org.apache.commons.lang3.StringUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import java.io.Serializable;
import java.lang.reflect.Method;
import java.util.Collections;
import java.util.List;

/**
 * 限流注解切面
 *
 * @author zhangshaojian
 * @date 2022/7/22 18:26
 */
@Aspect
@Configuration
public class LimitInterceptor {

    private static final Logger logger = LoggerFactory.getLogger(LimitInterceptor.class);
    private static final String UNKNOWN = "unknown";
    @Resource
    private RedisTemplate<String, Serializable> redisTemplate;

    @Around("execution(public * *(..)) && @annotation(com.dylan.redislimiting.annotation.RateLimiting)")
    public Object interceptor(ProceedingJoinPoint pjp) {
        MethodSignature signature = (MethodSignature) pjp.getSignature();
        Method method = signature.getMethod();
        RateLimiting limitAnnotation = method.getAnnotation(RateLimiting.class);
        LimitType limitType = limitAnnotation.limitType();
        String key;
        int limitPeriod = limitAnnotation.period();
        int limitCount = limitAnnotation.count();

        //根据限流类型获取不同的key，如果不传会以方法名作为key
        switch (limitType) {
            case IP:
                key = getIpAddress();
                break;
            case CUSTOMER:
                key = limitAnnotation.key();
                break;
            default:
                key = StringUtils.upperCase(method.getName());
        }
        String redisKey = StringUtils.isNoneBlank(limitAnnotation.prefix()) ? StringUtils.joinWith(":", limitAnnotation.prefix(), key) : key;
        try {
            boolean pass = limit(redisKey, limitCount, limitPeriod);
            if (pass) {
                return pjp.proceed();
            } else {
                throw new RuntimeException("亲,访问速度过快！稍后再试！！！");
            }
        } catch (Throwable e) {
            if (e instanceof RuntimeException) {
                throw new RuntimeException(e.getLocalizedMessage());
            }
            throw new RuntimeException("server exception");
        }
    }

    /**
     * 限流方法 true-放行；false-限流
     * key:缓存的key,limitCount:key的最大值,limitSecond:key的过期时间
     */
    public boolean limit(String key, int limitCount, int limitSecond) {
        List<String> keys = Collections.singletonList(key);
        String luaScript = buildLuaScript();
        RedisScript<Integer> redisScript = new DefaultRedisScript<>(luaScript, Integer.class);
        Integer count = redisTemplate.execute(redisScript, keys, limitCount, limitSecond);
        //判断是否超过数量=>放行or限流
        return count != null && count <= limitCount;
    }

    /**
     * 编写redis Lua 限流脚本
     * 流程：如果超过指定数量就返回，没有超过就计数+1，返回当前的值。第一次进来的时候还要给缓存设置过期时间。
     * c = redis.call('get',KEYS[1]) //从缓存中获取KEY
     * if c and tonumber(c) > tonumber(ARGV[1]) then  return c end //超过限定的最大值就直接返回
     * c = redis.call('incr',KEYS[1])    //没有超过计数器+1,返回当前c的值
     * if tonumber(c) == 1 then redis.call('expire',KEYS[1],ARGV[2]) end
     * return c;
     * 注:Incr命令将key中储存的数字值增一，返回当前key的数值。
     * 注:tonumber方法是Lua语法里面的函数，把参数转成10进制的数字
     * 注:KEYS[1] 表示key，ARGV[1]表示:value,ARGV[2]:超时时间
     */
    private String buildLuaScript() {
        StringBuilder lua = new StringBuilder();
        lua.append("local c");
        lua.append("\nc = redis.call('get',KEYS[1])");
        //实际调用次数超过阈值，则直接返回
        lua.append("\nif c and tonumber(c) > tonumber(ARGV[1]) then");
        lua.append("\nreturn tostring(c);");
        lua.append("\nend");
        // 执行计算器自加
        lua.append("\nc = redis.call('incr',KEYS[1])");
        lua.append("\nif tonumber(c) == 1 then");
        //从第一次调用开始限流，设置对应键值的过期
        lua.append("\nredis.call('expire',KEYS[1],ARGV[2])");
        lua.append("\nend");
        lua.append("\nreturn tostring(c);");
        return lua.toString();
    }

    //获取id地址
    public String getIpAddress() {
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.currentRequestAttributes()).getRequest();
        String ip = request.getHeader("x-forwarded-for");
        if (ip == null || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {
            ip = request.getHeader("Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {
            ip = request.getHeader("WL-Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {
            ip = request.getRemoteAddr();
        }
        return ip;
    }
}
